package cn.pengpeng;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapreduce.*;
import org.apache.hadoop.util.ReflectionUtils;
import org.apache.log4j.Logger;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.sql.*;
import java.util.ArrayList;
import java.util.List;

public class MySQLInputFormat<V extends MySQLInputWritable> extends InputFormat<LongWritable, V> {

    private static final Logger LOG = Logger.getLogger(MySQLInputFormat.class);

    /** 配置 - 输入端数据库驱动类 */
    public static final String MYSQL_INPUT_DRIVER = "mysql.input.driver";
    /** 配置 - 输入端数据库URL */
    public static final String MYSQL_INPUT_URL = "mysql.input.url";
    /** 配置 - 输入端数据库用户名 */
    public static final String MYSQL_INPUT_USERNAME = "mysql.input.username";
    /** 配置 - 输入端数据库密码 */
    public static final String MYSQL_INPUT_PASSWORD = "mysql.input.password";

    /** 配置 - 查询总记录数语句 */
    public static final String MYSQL_INPUT_SELECT_COUNT_SQL = "mysql.input.select.count";
    /** 配置 - 查询语句 */
    public static final String MYSQL_INPUT_SELECT_RECORD_SQL = "mysql.input.select.record";
    /** 配置 - 每个数据分片包含的条数（默认 100 条） */
    public static final String MYSQL_INPUT_SPLIT_PRE_SIZE = "mysql.input.split.pre.size";
    /** 配置 - 读取数据类 */
    public static final String MYSQL_OUTPUT_VALUE_CLASS = "mysql.output.value.class";


    /**
     * 计算切片，决定map任务数量
     */
    @Override
    public List<InputSplit> getSplits(JobContext context)
            throws IOException, InterruptedException {

        Configuration conf = context.getConfiguration();
        Connection conn = null;
        Statement stmt = null;
        ResultSet rs = null;

        long recordCount = 0;

        try {
            conn = this.getConnection(conf);
            stmt = conn.createStatement();
            rs = stmt.executeQuery(conf.get(MYSQL_INPUT_SELECT_COUNT_SQL));
            if(rs.next())
                recordCount = rs.getLong(1);
        } catch (Exception e) {
            throw new IOException("查询数据总量失败", e);
        } finally {
            this.closeAutoCloseable(conn);
            this.closeAutoCloseable(stmt);
            this.closeAutoCloseable(rs);
        }

        List<InputSplit> splits = new ArrayList<InputSplit>();
        // 计算分片数量
        long preSplitCount = conf.getLong(MYSQL_INPUT_SPLIT_PRE_SIZE, 100);
        int splitNums = (int) (recordCount / preSplitCount +
                recordCount % preSplitCount == 0 ? 0 : 1);

        // 将数据分片信息存入列表中
        for(int i = 0; i < splitNums; i++) {
            if(i != splitNums - 1)
                splits.add(new MySQLInputSplit(i * preSplitCount, (i + 1) * preSplitCount));
            else
                splits.add(new MySQLInputSplit(i * preSplitCount, recordCount));
        }

        return splits;
    }

    /**
     * 创建记录读取对象
     */
    @Override
    public RecordReader<LongWritable, V> createRecordReader(InputSplit split, TaskAttemptContext context)
            throws IOException, InterruptedException {

        RecordReader<LongWritable, V> reader = new MySQLRecordReader();
        reader.initialize(split, context);
        return reader;
    }


    /**
     * 获取数据库连接
     * @param conf
     * @return
     * @throws Exception
     */
    private Connection getConnection(Configuration conf) throws Exception {

        String driver = conf.get(MYSQL_INPUT_DRIVER);
        String url = conf.get(MYSQL_INPUT_URL);
        String username = conf.get(MYSQL_INPUT_USERNAME);
        String password = conf.get(MYSQL_INPUT_PASSWORD);

        Class.forName(driver);
        return DriverManager.getConnection(url, username, password);
    }

    /**
     * 关闭连接
     * @param autoCloseable
     */
    private void closeAutoCloseable(AutoCloseable autoCloseable) {
        try {
            if(autoCloseable != null)
                autoCloseable.close();
        } catch (Exception e) {
            LOG.error("关闭失败"+e.getMessage());
        }
    }

    /**
     * MySQL数据切片信息类
     */
    public static class MySQLInputSplit extends InputSplit implements Writable {
        // 分片数据位置信息，MySQL数据不存在HDFS中，所以数组设置为空
        private String[] locations = new String[0];
        // 开始位置
        private long start;
        // 结束位置
        private long end;

        public MySQLInputSplit() {
        }

        public MySQLInputSplit(long start, long end) {
            this.start = start;
            this.end = end;
        }

        @Override
        public long getLength() throws IOException, InterruptedException {
            return this.end - this.start;
        }

        @Override
        public String[] getLocations() throws IOException, InterruptedException {
            // 根据该值决定是否采用数据本地化策略
            return this.locations;
        }

        public long getStart() {
            return start;
        }

        public void setStart(long start) {
            this.start = start;
        }

        public long getEnd() {
            return end;
        }

        public void setEnd(long end) {
            this.end = end;
        }

        @Override
        public void write(DataOutput out) throws IOException {
            out.writeLong(this.start);
            out.writeLong(this.end);
        }

        @Override
        public void readFields(DataInput in) throws IOException {
            this.start = in.readLong();
            this.end = in.readLong();
        }
    }

    /**
     * MySQL数据读取类
     *
     * @param <V>
     */
    public class MySQLRecordReader extends RecordReader<LongWritable, V> {

        private Connection conn;
        private ResultSet rs = null;
        private Configuration conf;
        private MySQLInputSplit split;
        private LongWritable key = null;
        private V value = null;
        private long postion = 0; // 计算当前进度

        @Override
        public void initialize(InputSplit split, TaskAttemptContext context)
                throws IOException, InterruptedException {
            this.split = (MySQLInputSplit) split;
            this.conf = context.getConfiguration();
        }

        /**
         * 通过反射实例化输出类
         * 默认为空数据类型
         * @return
         */
        @SuppressWarnings("unchecked")
        private V createValue() {
            Class<? extends MySQLInputWritable> clazz = this.conf.getClass(MYSQL_OUTPUT_VALUE_CLASS,
                    MySQLNullWritable.class, MySQLInputWritable.class);
            return (V) ReflectionUtils.newInstance(clazz, this.conf);
        }

        /**
         * 组装查询语句
         * @return
         */
        private String getQuerySql() {
            String sql = conf.get(MYSQL_INPUT_SELECT_RECORD_SQL);
            try {
                sql += " LIMIT " + this.split.getLength();
                sql += " OFFSET " + this.split.getStart();
            } catch (Exception e) {
                LOG.error(e.getMessage());
            }

            return sql;
        }

        @Override
        public boolean nextKeyValue()
                throws IOException, InterruptedException {
            if(this.key == null) {
                this.key = new LongWritable();
            }
            if(this.value == null) {
                this.value = createValue();
            }
            if(this.conn == null) {
                try {
                    this.conn = MySQLInputFormat.this.getConnection(this.conf);
                } catch (Exception e) {
                    throw new IOException("获取数据库连接失败", e);
                }
            }

            try {
                if(this.rs == null) {
                    String sql = this.getQuerySql();
                    Statement stmt = this.conn.createStatement();
                    this.rs = stmt.executeQuery(sql);
                }

                if(!this.rs.next()) {
                    return false; // 没有下一个结果了
                }

                // 还有结果
                this.value.readFieldsFromResultSet(this.rs); // 读取字段信息
                this.key.set(this.postion);
                this.postion++; // 更新进度
                return true;
            } catch (SQLException e) {
                throw new IOException("获取数据失败", e);
            }
        }

        @Override
        public LongWritable getCurrentKey()
                throws IOException, InterruptedException {

            return this.key;
        }

        @Override
        public V getCurrentValue()
                throws IOException, InterruptedException {

            return this.value;
        }

        @Override
        public float getProgress()
                throws IOException, InterruptedException {

            return this.postion / this.split.getLength();
        }

        @Override
        public void close() throws IOException {
            MySQLInputFormat.this.closeAutoCloseable(this.conn);
            MySQLInputFormat.this.closeAutoCloseable(this.rs);
        }
    }

    /**
     * 空数据类型
     */
    public class MySQLNullWritable extends MySQLInputWritable {

        @Override
        public void write(DataOutput out) throws IOException {

        }

        @Override
        public void readFields(DataInput in) throws IOException {

        }

        @Override
        public void readFieldsFromResultSet(ResultSet rs) throws SQLException {

        }
    }
}